import pickle
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go


def read_pickle_file(filename):
  with open(filename, 'rb') as file:
    return pickle.load(file)


def inject_dict_into_dataframe(meta, data):
  meta_len = data.__len__()
  for idx, (key, value) in enumerate(meta.items()):
    tmp_row = value
    layers, kernel_size, lr = [int(str) for str in key.split('_')]
    tmp_row['num_layers'] = layers
    tmp_row['kernel_size'] = kernel_size * 2 + 1
    tmp_row['lr'] = lr
    data.loc[meta_len + idx] = tmp_row
  return data


data = None
for i in range(1, 9):
  meta = read_pickle_file(f'result{i}.pkl')

  if data is None:
    columns = []
    columns = ['num_layers', 'kernel_size', 'lr']
    columns.extend(meta[list(meta.keys())[0]].keys())
    data = pd.DataFrame(columns=columns)

  data = inject_dict_into_dataframe(meta, data)

fig = px.box(data, x='kernel_size', y='f1_Overall', points='all')
fig.show()

fig = px.box(data, x='num_layers', y='f1_Overall', points='all')
fig.show()

fig = px.box(data, x='kernel_size', y='acc_Overall', points='all')
fig.show()

fig = px.box(data, x='num_layers', y='acc_Overall', points='all')
fig.show()

fig = go.Figure()
for kernel_size in range(3, 20, 2):
  fig.add_trace(
      go.Box(
          x=data[data['kernel_size'] == kernel_size]['num_layers'],
          y=data[data['kernel_size'] == kernel_size]['f1_Overall'],
          name=f'kernel_size_{kernel_size}'))
fig.update_layout(
    xaxis_title='num_layers',
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

fig = go.Figure()
for kernel_size in range(3, 20, 2):
  fig.add_trace(
      go.Box(
          x=data[data['kernel_size'] == kernel_size]['num_layers'],
          y=data[data['kernel_size'] == kernel_size]['acc_Overall'],
          name=f'kernel_size_{kernel_size}'))
fig.update_layout(
    xaxis_title='num_layers',
    yaxis_title='acc_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

print(meta)